from langchain_community.embeddings import OpenAIEmbeddings
from langchain_core.embeddings import Embeddings
from pydantic import BaseModel
from ..config.agents import LLMType
from ..config.env import get_model_config
import os

class EmbConfig(BaseModel):
    model: str
    base_url: str | None = None
    api_key: str | None = None
    dimensions: int = 1536
    kwargs: dict = {}

    @classmethod
    def default(cls) -> "EmbConfig":
        model, base_url, api_key, kwargs = get_model_config(name=os.getenv(LLMType.LLM_EMB))
        return EmbConfig(
            model=model,
            base_url=base_url,
            api_key=api_key,
            dimensions=kwargs.pop("dimensions", 1536),
            kwargs=kwargs
        )

    @classmethod
    def build(cls,
              model: str,
              base_url: str | None = None,
              api_key: str | None = None,
              dimensions: int = 1536,
              kwargs: dict = None) -> "EmbConfig":
        return EmbConfig(
            model=model,
            base_url=base_url,
            api_key=api_key,
            dimensions=dimensions,
            kwargs=kwargs or {}
        )


def create_openai_emb(config: EmbConfig, **kwargs) -> Embeddings:
    combined_kwargs = {**config.kwargs, **kwargs}

    llm_kwargs = {"model": config.model, "dimensions": config.dimensions}

    if config.base_url:
        llm_kwargs["base_url"] = config.base_url

    if config.api_key:
        llm_kwargs["api_key"] = config.api_key

    llm_kwargs.update(combined_kwargs)

    emb = OpenAIEmbeddings(**llm_kwargs)
    return emb
